#pragma once

#include <vector>
#include <cstdint>
#include <string>
#include "core.h"


namespace uzu
{
    // fixed size tensor.
    class Tensor
    {
    private:
        float* m_Data = nullptr;
        float* m_DataGpu = nullptr;
        std::vector<uint32_t> m_Shape;
        DeviceType m_DeviceType = DeviceType::CPU;
    public:
        Tensor() { }
        Tensor(uint32_t size);
        Tensor(uint32_t width, uint32_t height);
        Tensor(uint32_t width, uint32_t height, uint32_t channels);
        ~Tensor();

        void CopyFromHost(float* data);
        uint32_t NDim() const { return (uint32_t)(m_Shape.size()); }
        std::vector<uint32_t> Shape() const { return m_Shape; }
        uint32_t Size();
        float* Data() const { return m_Data; }
        float* DataGpu() const { return m_DataGpu; }
        void Create(const std::vector<uint32_t>& shape, DeviceType type = DeviceType::CPU);

        void ToHost();
        void ToDevice();
        
    private:
        void Alloc(const std::vector<uint32_t>& shape);
        void AllocGpu(const std::vector<uint32_t>& shape);
        uint32_t ShapeSize(const std::vector<uint32_t>& shape);
    };
}
